package com.sailing.xjpb.elasticsearch.repository;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.Refresh;
import co.elastic.clients.elasticsearch.core.*;
import co.elastic.clients.elasticsearch.core.bulk.BulkOperation;
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.elasticsearch.core.search.HitsMetadata;
import com.sailing.xjpb.common.exception.EsException;
import com.sailing.xjpb.elasticsearch.EsClientPool;
import com.sailing.xjpb.elasticsearch.util.EsUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageImpl;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * es查询API
 *
 * @author luozy
 * createTime 2020/5/25
 */
@Slf4j
public class EsRepository {

    @Resource
    private EsClientPool esClientPool;

    /**
     * Description: 查询
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param searchRequest 查询
     * @param type          类型
     * @return List<T>
     */
    public <T> List<T> query(SearchRequest searchRequest, Class<T> type) {
        if (searchRequest == null) {
            throw new EsException("查询条件Search不能为空！");
        }
        if (type == null) {
            throw new EsException("查询结果类型不能为空！");
        }
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            log.debug(EsUtil.printEsBySearchRequest(searchRequest));
            SearchResponse<T> searchResponse = client.search(searchRequest, type);
            return this.list(searchResponse);
        } catch (Exception e) {
            throw new EsException("es查询失败！", e);
        }
    }

    /**
     * Description: 查询
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param requestBuilder 查询
     * @param type           类型
     * @return Page<T> 分页数据
     */
    public <T> Page<T> page(PageRequestBuilder requestBuilder, Class<T> type) {
        if (requestBuilder == null) {
            throw new EsException("查询条件Search不能为空！");
        }
        if (type == null) {
            throw new EsException("查询结果类型不能为空！");
        }
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            SearchRequest searchRequest = requestBuilder.build();
            log.debug(EsUtil.printEsBySearchRequest(searchRequest));
            SearchResponse<T> searchResponse = client.search(searchRequest, type);
            return this.page(requestBuilder, searchResponse);
        } catch (Exception e) {
            throw new EsException("es分页查询失败！", e);
        }
    }

    /**
     * Description: 计数
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param countRequest 计数
     * @return long
     */
    public long count(CountRequest countRequest) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            CountResponse count = client.count(countRequest);
            return count.count();
        } catch (Exception e) {
            throw new EsException("es计数查询失败！", e);
        }
    }

    /**
     * Description: 新增信息
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param indexName 索引名称
     * @param data      数据
     * @return long
     */
    public <T> long save(String indexName, T data) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            IndexResponse indexResponse = client.index(x -> x.index(indexName).refresh(Refresh.True).document(data));
            return indexResponse.version();
        } catch (Exception e) {
            throw new EsException("es数据保存失败！", e);
        }
    }

    /**
     * Description: 新增信息 指定ID
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param indexName 索引名称
     * @param id        数据主键
     * @param data      数据
     * @return long
     */
    public <T> long save(String indexName, String id, T data) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            IndexResponse indexResponse = client.index(x -> x.index(indexName).refresh(Refresh.True).id(id).document(data));
            return indexResponse.version();
        } catch (Exception e) {
            throw new EsException("es数据保存失败！", e);
        }
    }

    /**
     * Description: 修改信息 指定ID
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param indexName 索引名称
     * @param id        数据主键
     * @param data      数据
     * @return long
     */
    public <T> long update(String indexName, String id, T data) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            UpdateResponse<Object> userUpdateResponse = client.update(x -> x.index(indexName).refresh(Refresh.True).id(id).doc(data),
                    Object.class);
            return userUpdateResponse.version();
        } catch (Exception e) {
            throw new EsException("es数据修改失败！", e);
        }
    }

    /**
     * Description: 批量新增信息
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param indexName 索引名称
     * @param dataList  数据
     * @return long
     */
    public <T> List<BulkResponseItem> bulkInsert(String indexName, List<T> dataList) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            //创建BulkOperation列表准备批量插入doc
            List<BulkOperation> bulkOperations = new ArrayList<>();
            //将user中id作为es id，也可不指定id es会自动生成id
            dataList.forEach(doc -> bulkOperations
                    .add(BulkOperation
                            .of(b -> b
                                    .index(c -> c
                                            .document(doc)
                                    )
                            )
                    )
            );
            BulkResponse bulk = client
                    .bulk(x -> x
                            .index(indexName)
                            .operations(bulkOperations)
                    );
            return bulk.items();
        } catch (Exception e) {
            throw new EsException("es数据批量保存失败！", e);
        }
    }

    /**
     * Description: 批量新增信息 指定ID
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param indexName 索引名称
     * @param idList    数据主键
     * @param dataList  数据
     * @return long
     */
    public <T> List<BulkResponseItem> bulkInsert(String indexName, List<String> idList, List<T> dataList) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            //创建BulkOperation列表准备批量插入doc
            List<BulkOperation> bulkOperations = new ArrayList<>();
            //将id作为es id，也可不指定id es会自动生成id
            for (int i = 0; i < dataList.size(); i++) {
                int finalI = i;
                bulkOperations.add(BulkOperation
                        .of(b -> b
                                .index(c -> c
                                        .id(idList.get(finalI))
                                        .document(dataList.get(finalI)
                                        )
                                )
                        )
                );
            }
            BulkResponse bulk = client
                    .bulk(x -> x
                            .index(indexName)
                            .operations(bulkOperations)
                    );
            return bulk.items();
        } catch (Exception e) {
            throw new EsException("es数据批量保存失败！", e);
        }
    }

    /**
     * Description: 批量删除信息
     * CreateUser: 种传民
     * CreateDate:
     * <p>
     * EditDescription:
     * EditUser: 种传民
     * EditDate:
     *
     * @param indexName 索引名称
     * @param idList    数据主键
     * @return long
     */
    public List<BulkResponseItem> bulkDelete(String indexName, List<String> idList) {
        ElasticsearchClient client;
        try {
            client = esClientPool.getClient();
            // 构建批量操作对象BulkOperation的集合
            List<BulkOperation> bulkOperations = new ArrayList<>();
            // 向集合中添加需要删除的文档id信息
            for (int i = 0; i < idList.size(); i++) {
                int finalI = i;
                bulkOperations.add(BulkOperation.of(b -> b
                        .delete((d -> d
                                .index(indexName)
                                .id(idList.get(finalI)
                                )
                        ))
                ));
            }
            // 调用客户端的bulk方法，并获取批量操作响应结果
            BulkResponse response = client
                    .bulk(e -> e
                            .index(indexName)
                            .operations(bulkOperations));
            return response.items();
        } catch (Exception e) {
            throw new EsException("es数据删除失败！", e);
        }
    }

    private <T> Page<T> page(PageRequestBuilder requestBuilder, SearchResponse<T> search) {
        HitsMetadata<T> hitsMetadata = search.hits();
        long total = hitsMetadata.total().value();
        List<T> content = hitsMetadata.hits().stream().map(Hit::source).collect(Collectors.toList());
        return new PageImpl<>(content, requestBuilder.getPageable(), total);
    }

    private <T> List<T> list(SearchResponse<T> search) {
        return search.hits().hits().stream().map(Hit::source).collect(Collectors.toList());
    }

}
